import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_grid

# 创建一个 CIFAR-10 测试集的 Dataset 对象，后续可通过 DataLoader 批量加载数据
dataset = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True,
    transform=torchvision.transforms.ToTensor()
)
# 创建一个 数据加载器（DataLoader），用于高效加载和批量处理数据集
data_loader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=False)


# 网络定义
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 6, 3, padding=0, stride=1)

    def forward(self, x):
        return self.conv1(x)


net = Net()
writer = SummaryWriter("logs")

step = 0
for data in data_loader:
    images, labels = data
    output = net(images)
    print("Input shape:", images.shape)  # [64,3,32,32]
    print("Output shape:", output.shape)  # [64,6,30,30]

    """
    这个地方会报错，因为images是一个四维张量，而要求的是一个三维张量
    这个错误是因为 writer.add_image() 方法要求输入的图像张量格式必须是 (C, H, W)（通道、高度、宽度），
    但传入的是 (64, 3, 32, 32)（批次、通道、高度、宽度）。
    """
    # writer.add_image("input", images, step)
    # writer.add_image("output", output, step)

    # 方法1：使用 make_grid 拼接图像
    # writer.add_image("input_grid", make_grid(images), step)
    # writer.add_image("output_grid", make_grid(output[:, :3]), step)  # 取前3通道

    # # 方法2：记录单张图像
    writer.add_image("input_single", images[0], step)
    writer.add_image("output_single", output[0][:3], step)  # 取前3通道

    # 方法3：
    # writer.add_image("input", images, step)
    # output = torch.reshape(output, (-1, 3, 30, 30))
    # writer.add_image("output_single", output[0], step)

    step += 1

writer.close()
